# -*- coding: utf-8 -*-
import socket, time
import sys
from logger import logger
#温柔的telnet类，一定要回显了主机名（命令下发成功后）才执行下一条

class TelnetSock(object):
    '''
            只针对太一负载均衡设备
    '''
    def __init__(self, destip, dport = 23, dutname = 'host', 
                 username = 'admin', password = 'adc.admin', sendInterval = 0.1, log_level = 'off', sock_timeout=10):
        self.dest = (destip, dport)
        self.username = username
        self.password = password
        self.dutname = dutname
        self.sendInterval = sendInterval
        self.log_level = log_level
        self.logger = logger(dutname, self.log_level)
        try:
            self.tsock = socket.create_connection(self.dest, timeout=sock_timeout)
        except:
            self.logger.warn("Can not connect to the DUT %s:%s !"%self.dest)
            sys.exit()
        #控制信令协商
        t = self._recv()
        self.logger.debug(t)
        self.tsock.send('\xff\xfb\x01')
        t = self._recv()
        self.logger.debug(t)
        self.tsock.send('\xff\xfb\x1f\xff\xfa\x1f\x00\x50'
                        '\x00\x1d\xff\xf0\xff\xfd\x01\xff\xfd\x03')  
        t = self._recv()
        self.logger.debug(t)
        self.login()
    def _recv(self, buffsize=1024):
        try:
            return self.tsock.recv(buffsize)
        except Exception, e:
            self.logger.error('Socket Error: %s'%e)
            sys.exit()
    def _send(self, sendstr):
        try:
            self.tsock.send(sendstr + '\n')
        except Exception, e:
            self.logger.error('Socket Error ')
            sys.exit()
    def login(self):
        self._send(self.username)
        self.recv_until('Password:')
        self._send(self.password)
        res = self._recv()
        if res.find('password error') != -1:
            self.logger.error(res)
            self.close()
            sys.exit()
        elif res == '\r\n':
            res = self._recv()
            self.logger.debug(res)
        if res.find('>')!=-1:
            self.dutname = res.strip('\r\n').strip('> ')
    def recv_until(self, matchStr):
        res = ''
        while self.tsock:
            data = self._recv()
            res += data
            if not data:
                break
            elif data.find('--More--')!=-1:
                self.tsock.send(' ')
                continue
            elif data.find(matchStr) != -1:
                break
            else:
                pass
        return res
    def send_commands(self, commands, intervalTime = 0):
        commandList = [l.strip() for l in commands.split('\n') if l]
        recvdata = ''
        for c in commandList:
            data = self.post(c)
            recvdata += data
            if intervalTime:
                time.sleep(intervalTime)
            else:
                time.sleep(self.sendInterval)
        return recvdata
    def post(self, command, intervalTime = 0):
        self._send(command)
        recvData = self.recv_until(self.dutname)
        if intervalTime:
            time.sleep(intervalTime)
        else:
            time.sleep(self.sendInterval)
        self.logger.info(recvData)
        return recvData
    def run_with_user_input(self):
        while self.tsock:
            try:
                userInput = raw_input('')
                if userInput == 'exit':
                    self.close()
                    return
            except KeyboardInterrupt:
                self.close()
                return
            self.post(userInput)  
    def close(self):
        self.tsock.close()
        self.logger.exit()

class TelnetSockSerial(object):
    '''只针对串口服务器'''
    def __init__(self, destip, dutname, username = 'admin', 
                 password = 'adc.admin', dport = 23, sendInterval = 0.1):
        self.dest = (destip, dport)
        self.username = username
        self.password = password
        self.sendInterval = sendInterval
        self.dutname = dutname
        self.ifPrint = 0
        try:
            self.tsock = socket.create_connection(self.dest)
        except:
            print "Can not connect to ",self.dest
            sys.exit()
        self.tsock.recv(1024)
        self.tsock.send('\xff\xfd\x01\xff\xfd\03')
        self.enter = '\x0d\x0a'
        #针对串口，发送此指令后，不同的地方就是回车键不是\n了，而是'\x0d\x0a'
    def login(self):
        self._send(self.username)
        while 1:
            data = self._recv()
            if not data:
                break
            elif data.find('Password:')!=-1:
                self._send(self.password)
                break
    def post(self, command):
        #此函数用来发送简单命令
        self._send(command)
        recvData = ''
        while 1:
            data = self._recv()
            recvData += data
            if not data:
                self.close()
                break
            elif data.find('--More--')!=-1:
                self.tsock.send(' ')
                continue
            elif data.find(self.dutname)!=-1:
                break
        if self.ifPrint:
            print recvData
        return recvData
    def _send(self, command):
        self.tsock.send(command + self.enter)
    def _recv(self, buffsize=1024):
        try:
            return self.tsock.recv(buffsize)
        except Exception, e:
            logger.error('Socket Error: %s'%e)
    def recv_data_until(self, match):
        res = ''
        while 1:
            data = self._recv()
            if not data:
                break
            elif data.find(match) != -1:
                res += data
                break
            else:
                res += data
        return res  
    def close(self):
        self.tsock.close()
        self.logger.exit()        
        
if __name__ == "__main__":
    d = TelnetSock('192.168.10.1', password = 'dtjjx!12', log_level = 'all')
    a = '''enable
    show cpu usa
    show mem usa
    show ip route
    show run vs'''
    d.send_commands(a)